from cryptography.hazmat.primitives.asymmetric import rsa, padding
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
import os
import base64
def encrypt_using_public_key(data: str, public_key_string: str, use_pkc_padding: bool = True) -> str:
# Format the public key string into PEM format
pem_key = f"-----BEGIN PUBLIC KEY-----\n{public_key_string}\n-----END PUBLIC KEY-----"
# Load the public key
public_key = serialization.load_pem_public_key(pem_key.encode('utf-8'))
# Choose padding scheme
if use_pkc_padding:
chosen_padding = padding.PKCS1v15()
else:
chosen_padding = padding.OAEP(
mgf=padding.MGF1(algorithm=hashes.SHA256()),
algorithm=hashes.SHA256(),
label=None
)
# Encrypt the data
encrypted_data = public_key.encrypt(
data.encode('utf-8'),
chosen_padding
)
return base64.b64encode(encrypted_data).decode('utf-8')
def decrypt_using_private_key(encrypted_data: str, private_key_string: str, use_pkc_padding: bool = True) -> str:
# Format the private key string into PEM format
pem_key = f"-----BEGIN PRIVATE KEY-----\n{private_key_string}\n-----END PRIVATE KEY-----"
# Load the private key
private_key = serialization.load_pem_private_key(pem_key.encode('utf-8'), password=None)
# Choose padding scheme
if use_pkc_padding:
chosen_padding = padding.PKCS1v15()
else:
chosen_padding = padding.OAEP(
mgf=padding.MGF1(algorithm=hashes.SHA256()),
algorithm=hashes.SHA256(),
label=None
)
# Decrypt the data
decrypted_data = private_key.decrypt(
base64.b64decode(encrypted_data),
chosen_padding
)
return decrypted_data.decode('utf-8')
# Represents the response structure for encrypted data
class EncryptedResponse:
encrypted_data: str
aes_properties: str
aes: bool
def __init__(self, encrypted_data: str, aes_properties: str, aes: bool):
self.encrypted_data = encrypted_data
self.aes_properties = aes_properties
self.aes = aes
def encrypt_with_aes(data: str, public_key: str, use_pkc_padding: bool = True) -> EncryptedResponse:
aes_key = os.urandom(32) # Generate a random 256-bit AES key
iv = os.urandom(16) # Generate a random 128-bit IV
# Create AES cipher
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv))
encryptor = cipher.encryptor()
# Pad the data to be a multiple of the block size
padded_data = data + (16 - len(data) % 16) * chr(16 - len(data) % 16)
encrypted_data = encryptor.update(padded_data.encode('utf-8')) + encryptor.finalize()
# Combine AES key and IV, then encode in base64
aes_properties = base64.b64encode(aes_key).decode('utf-8') + "." + base64.b64encode(iv).decode('utf-8')
# Encrypt the AES properties using the public key
base64_encrypted_data = base64.b64encode(encrypted_data).decode('utf-8')
encrypted_aes_properties = encrypt_using_public_key(aes_properties, public_key, use_pkc_padding)
return EncryptedResponse(base64_encrypted_data, encrypted_aes_properties, True)
def decrypt_with_aes(properties: str, encrypted_data: str, private_key: str, use_pkc_padding: bool = True) -> str:
# Decrypt the AES properties using the private key
decrypted_aes_properties_in_base64 = decrypt_using_private_key(properties, private_key, use_pkc_padding)
# Split the decrypted properties to get AES key and IV
array_of_iv_n_key = decrypted_aes_properties_in_base64.split(".")
aes_key = base64.b64decode(array_of_iv_n_key[0])
iv = base64.b64decode(array_of_iv_n_key[1])
# Create AES cipher for decryption
cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv))
decryptor = cipher.decryptor()
# Decrypt the data
decrypted_padded_data = decryptor.update(base64.b64decode(encrypted_data)) + decryptor.finalize()
# Remove padding
pad_length = decrypted_padded_data[-1]
decrypted_data = decrypted_padded_data[:-pad_length]
return decrypted_data.decode('utf-8')
def main():
# Example usage
public_key = "<PUBLIC_KEY_STRING>"
private_key = "<PRIVATE_KEY_STRING>"
data_to_encrypt = "{'id': 12345, 'name': 'John Doe', 'email': 'johedoe@example.com'}"
encrypted_data = encrypt_with_aes(data_to_encrypt, public_key, use_pkc_padding=False)
print("Encrypted Data (AES):", encrypted_data.encrypted_data, "\n")
print("AES Properties (Encrypted):", encrypted_data.aes_properties, "\n")
decrypted_data = decrypt_with_aes(encrypted_data.aes_properties, encrypted_data.encrypted_data, private_key, use_pkc_padding=False)
print("Decrypted Data (AES):", decrypted_data)
# OUTPUT:
# Encrypted Data (AES): <BASE64_ENCRYPTED_DATA>
#
# AES Properties (Encrypted): <BASE64_ENCRYPTED_AES_PROPERTIES>
#
# Decrypted Data (AES): {'id': 12345, 'name': 'John Doe', 'email': 'johndoe@example.com'}
if __name__ == "__main__":
main()